#!/usr/bin/env python
# -*- coding: UTF-8 -*-

'''
Copyright (C) 2009 dbzhang800@gmail.com
All rights reserved.
''' 

import ctypes
from ctypes import *
from ctypes.wintypes import *

NULL = 0
ANYSIZE_ARRAY = 1
INVALID_HANDLE_VALUE = -1
GENERIC_READ = 0x80000000
GENERIC_WRITE = 0x40000000
FILE_SHARE_READ = 0x00000001
FILE_SHARE_WRITE = 0x00000002
OPEN_EXISTING = 3
DIGCF_PRESENT = 0x00000002
DIGCF_ALLCLASSE = 0x00000004
DIGCF_DEVICEINTERFACE = 0x00000010
SPDRP_DEVICEDESC = 0
SPDRP_HARDWAREID = 1
SPDRP_DRIVER = 9
NO_ERROR = 0
ERROR_INSUFFICIENT_BUFFER = 122
ERROR_NO_MORE_ITEMS = 259
def CLT_CODE(DeviceType, Function, Method, Access):
    return (DeviceType<<16)|(Access<<14)|(Function<<2)|Method
USB_GET_NODE_CONNECTION_INFORMATION_EX = 274
METHOD_BUFFERED = 0
FILE_ANY_ACCESS = 0
FILE_DEVICE_UNKNOWN = 0x00000022
FILE_DEVICE_USB = FILE_DEVICE_UNKNOWN 
IOCTL_GET_HCD_DRIVERKEY_NAME = 0x220424
IOCTL_USB_GET_ROOT_HUB_NAME = 0x220408
IOCTL_USB_GET_NODE_INFORMATION = 0x220408
IOCTL_USB_GET_NODE_CONNECTION_INFORMATION = 0x22040C
IOCTL_USB_GET_NODE_CONNECTION_DRIVERKEY_NAME = 0x220420
IOCTL_USB_GET_DESCRIPTOR_FROM_NODE_CONNECTION = 0x220410
IOCTL_USB_GET_NODE_CONNECTION_INFORMATION_EX = CLT_CODE(FILE_DEVICE_USB,
        USB_GET_NODE_CONNECTION_INFORMATION_EX,
        METHOD_BUFFERED,
        FILE_ANY_ACCESS)
USB_STRING_DESCRIPTOR_TYPE = 3
MAXIMUM_USB_STRING_LENGTH = 255
NoDeviceConnected = 0
DeviceConnected = 1

kernel32 = windll.kernel32
CreateFile = kernel32.CreateFileW
CloseHandle = kernel32.CloseHandle
DeviceIoControl = kernel32.DeviceIoControl

setupapi = windll.setupapi
SetupDiGetClassDevs = setupapi.SetupDiGetClassDevsW
SetupDiEnumDeviceInfo = setupapi.SetupDiEnumDeviceInfo
SetupDiGetDeviceRegistryProperty = setupapi.SetupDiGetDeviceRegistryPropertyW
SetupDiDestroyDeviceInfoList = setupapi.SetupDiDestroyDeviceInfoList
SetupDiEnumDeviceInterfaces = setupapi.SetupDiEnumDeviceInterfaces
SetupDiGetDeviceInterfaceDetail = setupapi.SetupDiGetDeviceInterfaceDetailW
class GUID(Structure):
    _fields_ = [("Data1", c_ulong),
                ("Data2", c_ushort),
                ("Data3", c_ushort),
                ("Data4", c_ubyte * 8)]

class SP_DEVINFO_DATA(Structure):
    _pack_ = 1
    _fields_ = [("cbSize", DWORD),
            ("ClassGuid", GUID),
            ("DevInst", DWORD),
            ("Reserved", c_void_p)
            ]

class SP_DEVICE_INTERFACE_DATA(Structure):
    _pack_ = 1
    _fields_ = [("cbSize", DWORD),
            ("InterfaceClassGuid", GUID),
            ("Flags", DWORD),
            ("Reserved", c_void_p)
            ]

class SP_DEVICE_INTERFACE_DETAIL_DATA(Structure):
    _pack_ = 1
    _fields_ = [("cbSize", DWORD),
            ("DevicePath", c_wchar*ANYSIZE_ARRAY)]

class USB_HCD_DRIVERKEY_NAME(Structure):
    _pack_ = 1
    _fields_ = [("ActualLength", ULONG),
            ("DriverKeyName", c_wchar*1)
            ]

class USB_ROOT_HUB_NAME(Structure):
    _pack_ = 1
    _fields_ = [("ActualLength", ULONG),
            ("RootHubName", c_wchar*1)
            ]

class USB_HUB_DESCRIPTOR(Structure):
    _pack_ = 1
    _fields_ = [("bDescriptorLength", c_ubyte),
            ("bDescriptorType", c_ubyte),
            ("bNumberOfPorts", c_ubyte),
            ("wHubCharacteristics", c_ushort),
            ("bPowerOnToPowerGood", c_ubyte),
            ("bHubControlCurrent", c_ubyte),
            ("bRemoveAndPowerMask", c_ubyte*64)]

class USB_HUB_INFORMATION(Structure):
    _pack_ = 1
    _fields_ = [("HubDescriptor", USB_HUB_DESCRIPTOR),
            ("HubIsBusPowered", c_ubyte)]

class USB_MI_PARENT_INFORMATION(Structure):
    _fields_ = [("NumberOfInterfaces", ULONG)]

class USB_NODE_INFORMATION(Structure):
    class MyU(Union):
        _pack_ = 1
        _fields_ = [("HubInformation", USB_HUB_INFORMATION),
                ("MiParentInformation", USB_MI_PARENT_INFORMATION)]
    _pack_ = 1
    _fields_ = [("NodeType", c_uint),
            ("u", MyU)]

class USB_DEVICE_DESCRIPTOR(Structure):
    _pack_ = 1
    _fields_ = [("bLength", c_ubyte),
            ("bDescriptor", c_ubyte),
            ("bcdUSB", c_ushort),
            ("bDeviceClass", c_ubyte),
            ("bDeviceSubClass", c_ubyte),
            ("bDeviceProtocol", c_ubyte),
            ("bMaxPacketSize0", c_ubyte),
            ("idVendor", c_ushort),
            ("idProduct", c_ushort),
            ("bcdDevice", c_ushort),
            ("iManufacturer", c_ubyte),
            ("iProduct", c_ubyte),
            ("iSerialNumber", c_ubyte),
            ("iNumConfiguration", c_ubyte),
            ]

class USB_ENDPOINT_DESCRIPTOR(Structure):
    _pack_ = 1
    _fields_ = [("bLength", c_ubyte),
            ("bDescriptorType", c_ubyte),
            ("bEndpointAddress", c_ubyte),
            ("bmAttributes", c_ubyte),
            ("wMaxPacketSize", c_ushort),
            ("bInterval", c_ubyte),
            ]

class USB_PIPE_INFO(Structure):
    _pack_ = 1
    _fields_ = [("EndpointDescriptor", USB_ENDPOINT_DESCRIPTOR),
            ("ScheduleOffset", ULONG)]

class USB_NODE_CONNECTION_INFORMATION_EX(Structure):
    _pack_ = 1
    _fields_ = [("ConnectionIndex", ULONG),
            ("DeviceDescriptor", USB_DEVICE_DESCRIPTOR),
            ("CurrentConfigurationValue", c_ubyte),
            ("Speed", c_ubyte),
            ("DeviceIsHub", c_ubyte),
            ("DeviceAddress", c_ushort),
            ("NumberOfOpenPipes", c_ulong),
            ("ConnectionStatus", c_uint),#enum
            ("PipeList", USB_PIPE_INFO)
            ]

class USB_NODE_CONNECTION_DRIVERKEY_NAME(Structure):
    _pack_ = 1 # without this, size will be 12
    _fields_ = [("ConnectionIndex", ULONG),
            ("ActualLength", ULONG),
            ("DriverKeyName", c_wchar*1)
            ]

GUID_DEVCLASS_DISPLAY = GUID(0x4d36e968L, 0xe325, 0x11ce, 
        (c_ubyte*8)(0xbf, 0xc1, 0x08, 0x00, 0x2b, 0xe1, 0x03, 0x18))
GUID_CLASS_USB_HOST_CONTROLLER = GUID(0x3ABF6F2D, 0x71C4, 0x462A, 
        (c_ubyte*8)(0x8A, 0x92, 0x1E, 0x68, 0x61, 0xE6, 0xAF, 0x27))
GUID_CLASS_USB_DEVICE = GUID(0xA5DCBF10L, 0x6530, 0x11D2, 
        (c_ubyte*8)(0x90, 0x1F, 0x00, 0xC0, 0x4F, 0xB9, 0x51, 0xED))


def getRootHubName(handle):
    bytesReturned = DWORD()
    hubName = USB_ROOT_HUB_NAME()
    res = DeviceIoControl(
            handle,
            IOCTL_USB_GET_ROOT_HUB_NAME,
            NULL,
            0,
            byref(hubName),
            6,
            byref(bytesReturned),
            NULL)
    if not res:
        print FormatError()
        return
    ctypes.resize(hubName, hubName.ActualLength)
    res = DeviceIoControl(
            handle,
            IOCTL_USB_GET_ROOT_HUB_NAME,
            NULL,
            0,
            byref(hubName),
            hubName.ActualLength,
            byref(bytesReturned),
            NULL)
    return wstring_at(byref(hubName, 4))

def getHCDDriverKeyName(handle):
    bytesReturned = DWORD()
    keyName = USB_HCD_DRIVERKEY_NAME()
    res = DeviceIoControl(
            handle,
            IOCTL_GET_HCD_DRIVERKEY_NAME,
            NULL,
            0,
            byref(keyName),
            6,
            byref(bytesReturned),
            NULL)
    if not res:
        print FormatError()
        return
    ctypes.resize(keyName, keyName.ActualLength)
    res = DeviceIoControl(
            handle,
            IOCTL_GET_HCD_DRIVERKEY_NAME,
            NULL,
            0,
            byref(keyName),
            keyName.ActualLength,
            byref(bytesReturned),
            NULL)
    #print wstring_at(byref(keyName, 4))
    return wstring_at(byref(keyName, 4))

def getDriverKeyName(handle, index):
    driverKeyName = USB_NODE_CONNECTION_DRIVERKEY_NAME()
    driverKeyName.ConnectionIndex = index
    bytesReturned = DWORD()
    res = DeviceIoControl(
                handle,
                IOCTL_USB_GET_NODE_CONNECTION_DRIVERKEY_NAME,
                byref(driverKeyName),
                sizeof(USB_NODE_CONNECTION_DRIVERKEY_NAME),
                byref(driverKeyName),
                sizeof(USB_NODE_CONNECTION_DRIVERKEY_NAME),
                byref(bytesReturned),
                None)
    ctypes.resize(driverKeyName, driverKeyName.ActualLength)
    res = DeviceIoControl(
                handle,
                IOCTL_USB_GET_NODE_CONNECTION_DRIVERKEY_NAME,
                byref(driverKeyName),
                sizeof(USB_NODE_CONNECTION_DRIVERKEY_NAME),
                byref(driverKeyName),
                driverKeyName.ActualLength,
                byref(bytesReturned),
                None)
    return wstring_at(byref(driverKeyName, 8))

def enumHubs(rootHubName):
    devicePath = ur"\\.\{}".format(rootHubName)
    hubHandle = CreateFile(
            devicePath,
            GENERIC_WRITE,
            FILE_SHARE_WRITE,
            NULL,
            OPEN_EXISTING,
            0,
            NULL)
    if hubHandle == INVALID_HANDLE_VALUE:
        print FormatError()
        exit(1)
    hubInfo = USB_NODE_INFORMATION()
    #print sizeof(USB_NODE_INFORMATION)
    bytesReturned = DWORD()
    DeviceIoControl(hubHandle,
            IOCTL_USB_GET_NODE_INFORMATION,
            NULL,
            0,
            byref(hubInfo),
            sizeof(USB_NODE_INFORMATION),
            byref(bytesReturned),
            NULL)
    if hubInfo.u.HubInformation.HubIsBusPowered:
        print "    Hub Power: Bus Power"
    else:
        print "    Hub Power: Self Power"
    hubCharacter = hubInfo.u.HubInformation.HubDescriptor.wHubCharacteristics
    if (hubCharacter & 0x0003) == 0x0000:
        print "    Power switching: Ganged"
    elif (hubCharacter & 0x0003) == 0x0001:
        print "    Power switching: Individual"
    else:
        print "    Power switching: None"
    if (hubCharacter & 0x0004) == 0x0000:
        print "    Compound device: None"
    else:
        print "    Compound device: Yes"
    if (hubCharacter & 0x0018) == 0x0000:
        print "    Over-current Protection: Global"
    elif (hubCharacter & 0x0018) == 0x0008:
        print "    Over-current Protection: Individual"
    else:
        print "    No Over-current Protection"
    print ""
    for index in range(1, hubInfo.u.HubInformation.HubDescriptor.bNumberOfPorts+1):
        getHubPort(hubHandle, index)
    CloseHandle(hubHandle)

def getHubPort(handle, portIndex):
    connectInfo = USB_NODE_CONNECTION_INFORMATION_EX()
    connectInfo.ConnectionIndex = portIndex
    bytesReturned = DWORD()
    res = DeviceIoControl(handle,
            IOCTL_USB_GET_NODE_CONNECTION_INFORMATION_EX,
            byref(connectInfo),
            sizeof(USB_NODE_CONNECTION_INFORMATION_EX),
            byref(connectInfo),
            sizeof(USB_NODE_CONNECTION_INFORMATION_EX),
            byref(bytesReturned),
            NULL)
    if not res:
        print FormatError()
    if connectInfo.ConnectionStatus == NoDeviceConnected:
        print "    [Port{}]NoDeviceConnected".format(portIndex)
    else:
        driverKeyName = getDriverKeyName(handle, portIndex)
        #print "    [Port{}]{}".format(portIndex, driverKeyName)
        portLabel = getUsbDeviceLabel(driverKeyName)
        print u"    [Port{}]{}".format(portIndex, portLabel)

def getUsbDeviceLabel(driverKeyName):
    hDevInfo = SetupDiGetClassDevs(byref(GUID_CLASS_USB_DEVICE),
            0,
            0,
            DIGCF_PRESENT | DIGCF_DEVICEINTERFACE)
    if hDevInfo == INVALID_HANDLE_VALUE:
        print GetLastError(), FormatError()
        return
    
    for i in range(10000):
        deviceInfoData = SP_DEVINFO_DATA()
        deviceInfoData.cbSize = sizeof(SP_DEVINFO_DATA)
        if SetupDiEnumDeviceInfo(hDevInfo, i, byref(deviceInfoData)):
            DataT = DWORD()
            buff = create_unicode_buffer(0)
            buff_lenth = DWORD(0)
            while not SetupDiGetDeviceRegistryProperty(
                    hDevInfo,
                    byref(deviceInfoData),
                    SPDRP_DRIVER,
                    byref(DataT),
                    byref(buff),
                    buff_lenth,
                    byref(buff_lenth)):
                if GetLastError() == ERROR_INSUFFICIENT_BUFFER:
                    buff = create_unicode_buffer(buff_lenth.value)
                else:
                    print FormatError()
                    break
            if  driverKeyName != buff.value:
                continue
            while not SetupDiGetDeviceRegistryProperty(
                    hDevInfo,
                    byref(deviceInfoData),
                    SPDRP_DEVICEDESC,
                    byref(DataT),
                    byref(buff),
                    buff_lenth,
                    byref(buff_lenth)):
                if GetLastError() == ERROR_INSUFFICIENT_BUFFER:
                    buff = create_unicode_buffer(buff_lenth.value)
                else:
                    print FormatError()
                    break
            return buff.value
    
if __name__ == '__main__':
    hDevInfo = SetupDiGetClassDevs(byref(GUID_CLASS_USB_HOST_CONTROLLER),
            0,
            0,
            DIGCF_PRESENT | DIGCF_DEVICEINTERFACE)
    
    if hDevInfo == INVALID_HANDLE_VALUE:
        print GetLastError(), FormatError()
        exit(1)
    
    # enum each host controller
    for i in range(10000):
        deviceInterfaceData = SP_DEVICE_INTERFACE_DATA()
        deviceInterfaceData.cbSize = sizeof(SP_DEVICE_INTERFACE_DATA)
        if not SetupDiEnumDeviceInterfaces(
                hDevInfo,
                NULL,
                byref(GUID_CLASS_USB_HOST_CONTROLLER),
                i,
                byref(deviceInterfaceData)):
            break
        # get description string of the host controller
        deviceInfoData = SP_DEVINFO_DATA()
        deviceInfoData.cbSize = sizeof(SP_DEVINFO_DATA)
        if SetupDiEnumDeviceInfo(hDevInfo, i, byref(deviceInfoData)):
            DataT = DWORD()
            buff = create_unicode_buffer(0)
            buff_lenth = DWORD(0)
            while not SetupDiGetDeviceRegistryProperty(
                    hDevInfo,
                    byref(deviceInfoData),
                    SPDRP_DEVICEDESC,
                    byref(DataT),
                    byref(buff),
                    buff_lenth,
                    byref(buff_lenth)):
                if GetLastError() == ERROR_INSUFFICIENT_BUFFER:
                    buff = create_unicode_buffer(buff_lenth.value)
                else:
                    print FormatError()
                    break
            print u"\n{}".format(buff.value)
            while not SetupDiGetDeviceRegistryProperty(
                    hDevInfo,
                    byref(deviceInfoData),
                    SPDRP_HARDWAREID,
                    byref(DataT),
                    byref(buff),
                    buff_lenth,
                    byref(buff_lenth)):
                if GetLastError() == ERROR_INSUFFICIENT_BUFFER:
                    buff = create_unicode_buffer(buff_lenth.value)
                else:
                    print FormatError()
                    break
            print u"Hardware ID: {}".format(buff.value) # like PCI\VEN_8086&DEV_283A&SUBSYS_022E1028&REV_02
    
        # get path of the host controller
        requestedSize = DWORD()
        SetupDiGetDeviceInterfaceDetail(
                hDevInfo,
                byref(deviceInterfaceData),
                NULL,
                0,
                byref(requestedSize),
                NULL)
        deviceDetailData = SP_DEVICE_INTERFACE_DETAIL_DATA()
        ctypes.resize(deviceDetailData, requestedSize.value)
        deviceDetailData.cbSize = ctypes.sizeof(SP_DEVICE_INTERFACE_DETAIL_DATA)
        SetupDiGetDeviceInterfaceDetail(
                hDevInfo,
                byref(deviceInterfaceData),
                byref(deviceDetailData),
                requestedSize.value,
                NULL,
                NULL)
        devicePath = wstring_at(byref(deviceDetailData, 4))
        print u"DevicePath:  {}".format(devicePath)

        # open host controller
        hostHandle = CreateFile(
                devicePath,
                GENERIC_WRITE,
                FILE_SHARE_WRITE,
                NULL,
                OPEN_EXISTING,
                0,
                NULL)
        if hostHandle == INVALID_HANDLE_VALUE:
            print FormatError()
            exit(1)
    
        print u"DriverKey: {}".format(getHCDDriverKeyName(hostHandle))
        rootHubName = getRootHubName(hostHandle)
        CloseHandle(hostHandle)
        print u"\nRootHub\n"
        print u"    Root Hub: {}".format(rootHubName)
        enumHubs(rootHubName)
    
    if GetLastError() != NO_ERROR and GetLastError() != ERROR_NO_MORE_ITEMS:
        print FormatError()
        exit(1)
        
    SetupDiDestroyDeviceInfoList(hDevInfo)

